/*
 * (C) Copyright 2013 George Daswani (georgedaswani@gmail.com) and others.
 *
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the GNU Lesser General Public License
 * (LGPL) version 2.1 which accompanies this distribution, and is available at
 * http://www.gnu.org/licenses/lgpl-2.1.html
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
 * Lesser General Public License for more details.
 *
 * Contributors:
 *    George Daswani (georgedaswani@gmail.com)
 */

package am.api.osgi;

import java.io.Serializable;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import javax.transaction.Status;
import javax.transaction.Synchronization;
import javax.transaction.SystemException;
import javax.transaction.Transaction;
import javax.transaction.TransactionManager;
import javax.transaction.xa.XAResource;

import org.apache.log4j.Logger;
import org.springframework.util.Assert;

import am.api.AMConnection;
import am.api.AMConnectionFactory;
import am.api.XAAMConnection;
import am.api.XAAMConnectionFactory;
import am.api.exception.AMConnectionException;
import am.api.model.AMCredential;
import am.api.util.NLS;
import am.api.wrapper.AMConnectionKey;
import am.api.wrapper.AMConnectionWrapper;

public class XAAMConnectionEnlistingWrapper implements AMConnectionFactory,
		Serializable {

	private class TransactionListener implements Synchronization {

		private final Object key;

		public TransactionListener(Object key) {
			this.key = key;
		}

		public void afterCompletion(int status) {

			AMConnection connection = connectionMap.remove(key);
			if (connection != null) {
				connection.close();
			}
		}

		public void beforeCompletion() {

		}

	}

	private final static Logger logger = Logger
			.getLogger(XAAMConnectionEnlistingWrapper.class);

	private final static long serialVersionUID = 1L;

	private transient Map<Object, AMConnection> connectionMap = new ConcurrentHashMap<Object, AMConnection>();

	private transient TransactionManager transactionManager;

	private XAAMConnectionFactory wrappedCF;

	public XAAMConnectionEnlistingWrapper() {
		super();
	}

	private void enlist(Transaction transaction, XAResource xaResource,
			Object key) throws AMConnectionException {
		try {
			transaction.enlistResource(xaResource);
			transaction.registerSynchronization(new TransactionListener(key));
		} catch (Exception e) {
			try {
				transactionManager.setRollbackOnly();
			} catch (IllegalStateException e1) {
				e1.printStackTrace();
			} catch (SystemException e1) {
				e1.printStackTrace();
			}
		}
	}

	@Override
	public AMConnection getConnection() throws AMConnectionException {

		Transaction transaction = getTransaction();

		logger.debug(String.format("transaction = [%s]", transaction));

		if (transaction != null) {

			Object key = new AMConnectionKey(null, transaction);

			logger.debug(String.format("key = [%s]", key));

			AMConnection connection = connectionMap.get(key);

			logger.debug(String.format("connection = [%s]", connection));

			if (connection == null) {
				XAAMConnection xaConnection = wrappedCF.getXAConnection();
				connection = xaConnection.getConnection();
				enlist(transaction, xaConnection.getXAResource(), key);
				connectionMap.put(key, connection);
			}
			return getEnlistedConnection(connection, true);
		} else {
			return getEnlistedConnection(wrappedCF.getXAConnection()
					.getConnection(), false);
		}
	}

	@Override
	public AMConnection getConnection(AMCredential credential)
			throws AMConnectionException {
		Transaction transaction = getTransaction();
		if (transaction != null) {
			Object key = new AMConnectionKey(credential, transaction);
			AMConnection connection = connectionMap.get(key);
			if (connection == null) {
				XAAMConnection xaConnection = wrappedCF
						.getXAConnection(credential);
				connection = xaConnection.getConnection();
				enlist(transaction, xaConnection.getXAResource(), key);
				connectionMap.put(key, connection);
			}
			return getEnlistedConnection(connection, true);
		} else {
			return getEnlistedConnection(wrappedCF.getXAConnection(credential)
					.getConnection(), false);
		}
	}

	private AMConnection getEnlistedConnection(AMConnection connection,
			boolean enlisted) throws AMConnectionException {

		AMConnectionWrapper wrapper = new AMConnectionWrapper(connection,
				enlisted);

		logger.debug(String.format("wrapper = [%s]", wrapper));

		return wrapper;
	}

	private Transaction getTransaction() throws AMConnectionException {
		try {
			return (transactionManager.getStatus() == Status.STATUS_ACTIVE) ? transactionManager
					.getTransaction() : null;
		} catch (SystemException e) {
			AMConnectionException amE = new AMConnectionException(
					NLS.ERRORS.getString("transaction.cannot.retrieve"));
			amE.initCause(e);
			throw amE;
		}
	}

	public TransactionManager getTransactionManager() {
		return transactionManager;
	}

	public XAAMConnectionFactory getWrappedCF() {
		return wrappedCF;
	}

	public void init() {
		Assert.notNull(transactionManager, "transactionManager is required.");
		Assert.notNull(wrappedCF, "wrappedCF is required.");
	}

	public void setTransactionManager(TransactionManager transactionManager) {
		this.transactionManager = transactionManager;
	}

	public void setWrappedCF(XAAMConnectionFactory wrappedCF) {
		this.wrappedCF = wrappedCF;
	}

}
